const express = require('express');
const WebSocket = require('ws');
const { Client } = require('ssh2');
const path = require('path');
const cookieParser = require('cookie-parser');
const jwt = require('jsonwebtoken');
const cors = require('cors');
const fs = require('fs');
const util = require('util');
const moment = require('moment');
const yargs = require('yargs/yargs');
const { hideBin } = require('yargs/helpers');

// 颜色定义
const colors = {
    reset: '\x1b[0m',
    bright: '\x1b[1m',
    dim: '\x1b[2m',
    underscore: '\x1b[4m',
    blink: '\x1b[5m',
    reverse: '\x1b[7m',
    hidden: '\x1b[8m',
    
    // 前景色
    fg: {
        black: '\x1b[30m',
        red: '\x1b[31m',
        green: '\x1b[32m',
        yellow: '\x1b[33m',
        blue: '\x1b[34m',
        magenta: '\x1b[35m',
        cyan: '\x1b[36m',
        white: '\x1b[37m',
    },
    
    // 背景色
    bg: {
        black: '\x1b[40m',
        red: '\x1b[41m',
        green: '\x1b[42m',
        yellow: '\x1b[43m',
        blue: '\x1b[44m',
        magenta: '\x1b[45m',
        cyan: '\x1b[46m',
        white: '\x1b[47m',
    }
};

// 解析命令行参数
const argv = yargs(hideBin(process.argv))
  .option('cfg', {
    alias: 'c',
    type: 'string',
    description: 'Path to configuration JSON file',
    default: null
  })
  .argv;

// 默认配置
let CONFIG = {
    JWT_SECRET: 'your-secret-key-here',
    USERS: {
        admin: {
            password: 'admin@2025',
            servers: {
                default: {
                    host: '127.0.0.1',
                    port: 22,
                    username: 'test',
                    password: 'test@passwd2025'
                }
            }
        }
    },
    PORT: 8080
};

// 深度合并对象函数
function deepMerge(target, source) {
    const result = { ...target };
    for (const key in source) {
        if (source[key] instanceof Object && key in target) {
            result[key] = deepMerge(target[key], source[key]);
        } else {
            result[key] = source[key];
        }
    }
    return result;
}

// 验证配置
function validateConfig(config) {
    if (!config.JWT_SECRET || config.JWT_SECRET.length < 16) {
        throw new Error('JWT_SECRET must be at least 16 characters');
    }
    
    if (!config.USERS || Object.keys(config.USERS).length === 0) {
        throw new Error('At least one user must be defined in config');
    }
    
    for (const [username, user] of Object.entries(config.USERS)) {
        if (!user.password) {
            throw new Error(`User ${username} is missing password`);
        }
        
        if (!user.servers || !user.servers.default) {
            throw new Error(`User ${username} is missing default server config`);
        }
        
        const server = user.servers.default;
        if (!server.host || !server.port || !server.username || !server.password) {
            throw new Error(`Invalid server config for user ${username}`);
        }
    }
}

// 从配置文件加载配置
if (argv.cfg) {
    try {
        const configPath = path.resolve(argv.cfg);
        console.log(`Loading config from: ${configPath}`);
        
        const fileContent = fs.readFileSync(configPath, 'utf8');
        const fileConfig = JSON.parse(fileContent);
        
        // 深度合并配置
        CONFIG = deepMerge(CONFIG, fileConfig);
        validateConfig(CONFIG);
        
        console.log('Configuration loaded successfully');
    } catch (err) {
        console.error(`Failed to load config file: ${err.message}`);
        process.exit(1);
    }
}

const app = express();
// 使用配置中的端口
const PORT = CONFIG.PORT;

// 确保日志目录存在
const LOG_DIR = process.env.LOG_DIR || path.join(process.cwd(), 'logs');
if (!fs.existsSync(LOG_DIR)) {
    try {
        fs.mkdirSync(LOG_DIR, { recursive: true });
    } catch (err) {
        console.error('无法创建日志目录:', err);
        process.exit(1);
    }
}

// 当前日志文件
const LOG_FILE = path.join(LOG_DIR, `webssh_${moment().format('YYYY-MM-DD_HH-mm-ss')}.log`);

// 日志函数
function logToFile(message) {
    const timestamp = moment().format('YYYY-MM-DD HH:mm:ss');
    const logMessage = `[${timestamp}] ${message}\n`;
    
    fs.appendFile(LOG_FILE, logMessage, (err) => {
        if (err) console.error('写入日志文件失败:', err);
    });
}

// 带颜色的控制台日志
function colorLog(message, color = 'fg.white', bgColor = null) {
    const timestamp = moment().format('YYYY-MM-DD HH:mm:ss');
    const coloredMessage = `${colors.fg.cyan}[${timestamp}]${colors.reset} ` + 
                          (bgColor ? colors[bgColor] : '') +
                          (color ? colors[color] : '') +
                          message + 
                          colors.reset;
    
    console.log(coloredMessage);
    logToFile(message);
}

// 启动横幅
function showBanner() {
    const banner = `
${colors.fg.green}╔╔══════════════════════════════════════════════════╗╗
║║${colors.fg.yellow}          WebSSH 服务器启动中...              ${colors.fg.green}║║
║║${colors.fg.cyan}          版本: 1.2.1 | 端口: ${PORT}           ${colors.fg.green}║║
╚╚══════════════════════════════════════════════════╝╝
${colors.reset}`;
    
    console.log(banner);
    logToFile('WebSSH 服务器启动...');
}

// 中间件
app.use(cors());
app.use(express.json());
app.use(cookieParser());
app.use(express.static(path.join(__dirname, 'public')));

// 处理根路径处理
const isPackaged = typeof process.pkg !== 'undefined';
const staticPath = isPackaged 
  ? path.join(process.cwd(), 'public')
  : path.join(__dirname, 'public');

app.use(express.static(staticPath));

app.get('/', (req, res) => {
  res.sendFile(path.join(staticPath, 'index.html'));
});

// 登录路由
app.post('/login', (req, res) => {
    const { username, password } = req.body;
    
    if (!username || !password) {
        colorLog(`登录失败: 用户名或密码为空`, 'fg.red');
        return res.status(400).json({ message: '用户名和密码不能为空' });
    }
    
    const user = CONFIG.USERS[username];
    
    if (!user || user.password !== password) {
        colorLog(`登录失败: 用户名或密码错误 (用户名: ${username})`, 'fg.red');
        return res.status(401).json({ message: '用户名或密码错误' });
    }
    
    const token = jwt.sign({ username }, CONFIG.JWT_SECRET, { expiresIn: '1h' });
    
    colorLog(`用户 ${username} 登录成功`, 'fg.green');
    res.json({ 
        success: true,
        token 
    });
});

// 验证JWT中间件
function authenticateToken(req, res, next) {
    const token = req.cookies.token || req.headers['authorization']?.split(' ')[1];
    
    if (!token) {
        colorLog('未授权的请求: 缺少token', 'fg.yellow');
        return res.sendStatus(401);
    }
    
    jwt.verify(token, CONFIG.JWT_SECRET, (err, user) => {
        if (err) {
            colorLog(`无效的token: ${err.message}`, 'fg.red');
            return res.sendStatus(403);
        }
        
        req.user = user;
        next();
    });
}

// 创建HTTP服务器
const server = app.listen(PORT, () => {
    showBanner();
    colorLog(`HTTP服务器已启动，监听端口 ${PORT}`, 'fg.green');
    colorLog(`当前用户数量: ${Object.keys(CONFIG.USERS).length}`, 'fg.cyan');
});

// 创建WebSocket服务器
const wss = new WebSocket.Server({ 
    server,
    verifyClient: (info, done) => {
        const cookies = info.req.headers.cookie;
        const token = cookies?.split(';')
            .find(c => c.trim().startsWith('token='))
            ?.split('=')[1];
            
        if (!token) {
            colorLog('WebSocket连接尝试: 缺少token', 'fg.yellow');
            done(false, 401, 'Unauthorized');
            return;
        }
        
        jwt.verify(token, CONFIG.JWT_SECRET, (err, decoded) => {
            if (err) {
                colorLog(`无效的WebSocket token: ${err.message}`, 'fg.red');
                done(false, 403, 'Invalid token');
            } else {
                info.req.user = decoded;
                colorLog(`WebSocket连接授权: 用户 ${decoded.username}`, 'fg.cyan');
                done(true);
            }
        });
    }
});

// 跟踪所有活跃的WebSocket连接
const activeConnections = new Set();

wss.on('connection', (ws, req) => {
    const { username } = req.user;
    const userConfig = CONFIG.USERS[username];
    
    // 添加到活跃连接集合
    activeConnections.add(ws);
    colorLog(`新的WebSocket连接: 用户 ${username} (当前连接数: ${activeConnections.size})`, 'fg.cyan');
    
    ws.on('message', (message) => {
        try {
            const data = JSON.parse(message);
            
            if (data.action === 'connect') {
                colorLog(`用户 ${username} 正在尝试SSH连接...`, 'fg.cyan');
                const serverConfig = userConfig.servers.default;
                
                const sshClient = new Client();
                // 将SSH客户端附加到WebSocket对象上
                ws.sshClient = sshClient;
                
                sshClient.on('ready', () => {
                    colorLog(`用户 ${username} SSH连接成功 (${serverConfig.host}:${serverConfig.port})`, 'fg.green');
                    ws.send(JSON.stringify({ type: 'output', data: '\r\nSSH连接成功\r\n'}));
                    
                    sshClient.shell((err, stream) => {
                        if (err) {
                            colorLog(`用户 ${username} 创建shell会话失败: ${err.message}`, 'fg.red');
                            ws.send(JSON.stringify({ 
                                type: 'error', 
                                message: '创建shell会话失败: ' + err.message 
                            }));
                            return;
                        }
                        
                        colorLog(`用户 ${username} shell会话已启动`, 'fg.green');
                        
                        stream.on('data', (data) => {
                            ws.send(JSON.stringify({ type: 'output', data: data.toString('utf-8') }));
                        }).on('close', () => {
                            colorLog(`用户 ${username} shell会话已关闭`, 'fg.yellow');
                            ws.send(JSON.stringify({ type: 'disconnect', message: 'SSH会话已关闭' }));
                            sshClient.end();
                        }).stderr.on('data', (data) => {
                            colorLog(`用户 ${username} SSH错误: ${data.toString('utf-8')}`, 'fg.red');
                            ws.send(JSON.stringify({ 
                                type: 'error', 
                                message: 'SSH错误: ' + data.toString('utf-8') 
                            }));
                        });
                        
                        ws.on('message', (msg) => {
                            try {
                                const msgData = JSON.parse(msg);
                                if (msgData.action === 'input' && stream.writable) {
                                    stream.write(msgData.data);
                                }
                            } catch (e) {
                                colorLog(`处理WebSocket消息错误: ${e.message}`, 'fg.red');
                            }
                        });
                    });
                }).on('error', (err) => {
                    colorLog(`用户 ${username} SSH连接错误: ${err.message}`, 'fg.red');
                    ws.send(JSON.stringify({ 
                        type: 'error', 
                        message: 'SSH连接错误: ' + err.message 
                    }));
                }).connect({
                    host: serverConfig.host,
                    port: serverConfig.port,
                    username: serverConfig.username,
                    password: serverConfig.password,
                    tryKeyboard: true
                });
            }
        } catch (err) {
            colorLog(`WebSocket消息处理错误: ${err.message}`, 'fg.red');
            ws.send(JSON.stringify({ 
                type: 'error', 
                message: '处理消息时出错: ' + err.message 
            }));
        }
    });
    
    ws.on('close', () => {
        colorLog(`WebSocket连接关闭: 用户 ${username} (剩余连接数: ${activeConnections.size - 1})`, 'fg.yellow');
        // 关闭关联的SSH连接
        if (ws.sshClient) {
            ws.sshClient.end();
        }
        // 从活跃连接集合中移除
        activeConnections.delete(ws);
    });
});

// 错误处理
process.on('uncaughtException', (err) => {
    colorLog(`未捕获的异常: ${err.stack}`, 'fg.red');
    logToFile(`未捕获的异常: ${err.stack}`);
});

process.on('unhandledRejection', (reason, promise) => {
    colorLog(`未处理的拒绝: ${util.inspect(reason)}`, 'fg.red');
    logToFile(`未处理的拒绝: ${util.inspect(reason)}`);
});

// 优雅关闭处理
function gracefulShutdown() {
    colorLog('开始关闭服务器...', 'fg.yellow');
    
    // 设置10秒后强制退出
    const forceExitTimer = setTimeout(() => {
        colorLog('强制退出...', 'fg.red');
        process.exit(1);
    }, 10000);
    
    // 1. 关闭所有WebSocket连接
    colorLog(`正在关闭 ${activeConnections.size} 个活跃连接...`, 'fg.yellow');
    activeConnections.forEach(ws => {
        if (ws.sshClient) {
            ws.sshClient.end();
        }
        if (ws.readyState === WebSocket.OPEN) {
            ws.close();
        }
    });
    
    // 2. 关闭WebSocket服务器
    wss.close(() => {
        colorLog('WebSocket服务器已关闭', 'fg.green');
        
        // 3. 关闭HTTP服务器
        server.close(() => {
            colorLog('HTTP服务器已关闭', 'fg.green');
            clearTimeout(forceExitTimer);
            process.exit(0);
        });
    });
}

// 处理各种退出信号
['SIGINT', 'SIGTERM', 'SIGQUIT'].forEach(signal => {
    process.on(signal, () => {
        colorLog(`收到 ${signal} 信号，正在关闭服务器...`, 'fg.yellow');
        gracefulShutdown();
    });
});